import numpy as np
import pandas as pd
import talib
import matplotlib.pyplot as plt
from matplotlib import gridspec

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题


class TrendTradingSystem:
    def __init__(self, data):
        """
        初始化趋向交易系统

        参数:
        data: pandas DataFrame, 包含开盘价(open), 最高价(high), 最低价(low), 收盘价(close), 成交量(volume)
        """
        self.data = data.copy()
        self.signals = pd.DataFrame(index=data.index)
        self.signals['signal'] = 0  # 交易信号: 1买入, -1卖出, 0无信号

    def calculate_indicators(self):
        """计算各类技术指标"""
        # 趋势指标
        self.data['ma10'] = talib.MA(self.data['close'], timeperiod=10)
        self.data['ma30'] = talib.MA(self.data['close'], timeperiod=30)
        self.data['ma60'] = talib.MA(self.data['close'], timeperiod=60)

        # MACD指标
        self.data['macd'], self.data['macd_signal'], self.data['macd_hist'] = talib.MACD(
            self.data['close'], fastperiod=12, slowperiod=26, signalperiod=9)

        # RSI指标
        self.data['rsi'] = talib.RSI(self.data['close'], timeperiod=14)

        # 布林带
        self.data['upper'], self.data['middle'], self.data['lower'] = talib.BBANDS(
            self.data['close'], timeperiod=20, nbdevup=2, nbdevdn=2, matype=0)

        # ADX平均趋向指数
        self.data['adx'] = talib.ADX(
            self.data['high'], self.data['low'], self.data['close'], timeperiod=14)

        # ATR平均真实波幅
        self.data['atr'] = talib.ATR(
            self.data['high'], self.data['low'], self.data['close'], timeperiod=14)

    def generate_signals(self):
        """生成交易信号"""
        # 趋势确认条件: 短期均线上穿长期均线 + ADX > 25(强趋势)
        self.signals['trend_condition'] = np.where(
            (self.data['ma10'] > self.data['ma30']) &
            (self.data['ma30'] > self.data['ma60']) &
            (self.data['adx'] > 25), 1, 0)

        # MACD信号: MACD线上穿信号线
        self.signals['macd_signal'] = np.where(
            (self.data['macd'] > self.data['macd_signal']) &
            (self.data['macd'].shift(1) <= self.data['macd_signal'].shift(1)), 1, 0)

        # RSI超卖信号
        self.signals['rsi_signal'] = np.where(self.data['rsi'] < 30, 1, 0)

        # 布林带下轨反弹信号
        self.signals['bollinger_signal'] = np.where(
            (self.data['close'] < self.data['lower']) &
            (self.data['close'].shift(1) >= self.data['lower'].shift(1)), 1, 0)

        # 综合买入信号(需要至少3个条件满足)
        self.signals['buy_signal'] = np.where(
            (self.signals['trend_condition'] +
             self.signals['macd_signal'] +
             self.signals['rsi_signal'] +
             self.signals['bollinger_signal']) >= 3, 1, 0)

        # 卖出信号: 趋势反转或止盈止损
        # 趋势反转信号
        self.signals['trend_reversal'] = np.where(
            (self.data['ma10'] < self.data['ma30']) &
            (self.data['adx'] < 20), 1, 0)

        # RSI超买信号
        self.signals['rsi_overbought'] = np.where(self.data['rsi'] > 70, 1, 0)

        # 布林带上轨回落信号
        self.signals['bollinger_exit'] = np.where(
            (self.data['close'] > self.data['upper']) &
            (self.data['close'].shift(1) <= self.data['upper'].shift(1)), 1, 0)

        # 综合卖出信号(需要至少2个条件满足)
        self.signals['sell_signal'] = np.where(
            (self.signals['trend_reversal'] +
             self.signals['rsi_overbought'] +
             self.signals['bollinger_exit']) >= 2, 1, 0)

        # 生成最终信号(1买入, -1卖出)
        self.signals['signal'] = self.signals['buy_signal'] - self.signals['sell_signal']

    def backtest(self, initial_capital=100000.0):
        """简单的回测系统"""
        positions = pd.DataFrame(index=self.signals.index).fillna(0.0)
        positions['position'] = self.signals['signal'].cumsum()

        # 创建投资组合
        portfolio = positions.multiply(self.data['close'], axis=0)
        pos_diff = positions.diff()

        # 加入现金持仓
        portfolio['holdings'] = positions['position'] * self.data['close']
        portfolio['cash'] = initial_capital - (pos_diff['position'] * self.data['close']).cumsum()
        portfolio['total'] = portfolio['cash'] + portfolio['holdings']
        portfolio['returns'] = portfolio['total'].pct_change()

        return portfolio

    def plot_results(self, portfolio):
        """可视化结果"""
        fig = plt.figure(figsize=(14, 10))
        gs = gridspec.GridSpec(3, 1, height_ratios=[3, 1, 1])

        # 价格和信号图
        ax0 = plt.subplot(gs[0])
        self.data['close'].plot(ax=ax0, label='价格', color='black')

        # 标记买入信号
        ax0.plot(self.signals.loc[self.signals['signal'] == 1.0].index,
                 self.data['close'][self.signals['signal'] == 1.0],
                 '^', markersize=10, color='g', label='买入信号')

        # 标记卖出信号
        ax0.plot(self.signals.loc[self.signals['signal'] == -1.0].index,
                 self.data['close'][self.signals['signal'] == -1.0],
                 'v', markersize=10, color='r', label='卖出信号')

        # 绘制均线
        self.data[['ma10', 'ma30', 'ma60']].plot(ax=ax0)

        ax0.set_title('趋向交易系统 - 价格与交易信号')
        ax0.set_ylabel('价格')
        ax0.legend()
        ax0.grid()

        # MACD指标图
        ax1 = plt.subplot(gs[1], sharex=ax0)
        self.data['macd'].plot(ax=ax1, color='b', label='MACD线')
        self.data['macd_signal'].plot(ax=ax1, color='g', label='信号线')

        # 绘制MACD柱状图
        ax1.bar(self.data.index, self.data['macd_hist'],
                color=np.where(self.data['macd_hist'] > 0, 'r', 'g'),
                label='MACD柱')

        ax1.set_title('MACD指标')
        ax1.set_ylabel('MACD')
        ax1.legend()
        ax1.grid()

        # 资金曲线图
        ax2 = plt.subplot(gs[2], sharex=ax0)
        portfolio['total'].plot(ax=ax2, label='账户总值', color='b')

        ax2.set_title('资金曲线')
        ax2.set_ylabel('资金')
        ax2.legend()
        ax2.grid()

        plt.tight_layout()
        plt.show()


# 示例使用
if __name__ == "__main__":
    # 获取示例数据(这里使用yahoo finance的数据)
    import pandas_datareader as pdr

    data = pdr.get_data_yahoo('AAPL', start='2020-01-01', end='2023-01-01')
    # 下载苹果公司股票数据
    # data = yf.download('AAPL', start='2020-01-01', end='2023-01-01')
    data = data[['Open', 'High', 'Low', 'Close', 'Volume']]
    data.columns = ['open', 'high', 'low', 'close', 'volume']

    # 创建交易系统实例
    system = TrendTradingSystem(data)

    # 计算技术指标
    system.calculate_indicators()

    # 生成交易信号
    system.generate_signals()

    # 运行回测
    portfolio = system.backtest()

    # 打印最终收益
    final_return = (portfolio['total'][-1] / 100000 - 1) * 100
    print(f"最终收益率: {final_return:.2f}%")

    # 可视化结果
    system.plot_results(portfolio)